import chromadb
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.vector_store.chroma import ChromaVectorStore

# 定义向量存储数据库
chroma_client = chromadb.PersistentClient()

# 尝试获取集合，如果不存在则创建
try:
    chroma_collection = chroma_client.get_collection('quickstart')
    print('使用已经存在的本地知识库')
except chromadb.errors.InvalidCollectionException:
    chroma_collection = chroma_client.create_collection('quickstart')
    print('创建一个全新的本地知识库')

# 声明向量存储
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# 初始化一个HuggingFaceEmbedding对象，用于将文本转换为向量表示
embed_model = HuggingFaceEmbedding(
    # 指定了一个预训练的sentence-transformer 模型路径
    model_name='E:\AIModels\paraphrase\Ceceliachenen\paraphrase-multilingual-MiniLM-L12-v2'
)
# 将创建的嵌入模型赋值给全局设置的embed_model属性，这样在后续的索引构建过程中，就会使用这个模型
Settings.embed_model = embed_model

llm = HuggingFaceLLM(
    model_name='E:\AIModels\deepseek-ai\DeepSeek-R1-Distill-Qwen-1___5B',
    tokenizer_name='E:\AIModels\deepseek-ai\DeepSeek-R1-Distill-Qwen-1___5B',
    model_kwargs={'trust_remote_code': True},
    tokenizer_kwargs={'trust_remote_code': True}
)

# 设置全局的llm属性，这样在索引查询时会使用这个模型
Settings.llm = llm

# 从指定目录读取文档，将数据加载到内存
documents = SimpleDirectoryReader(r'data').load_data()

# 创建节点解析器
node_parser = SimpleNodeParser.from_defaults(chunk_size=512)
# 将文档分割成节点
base_node = node_parser.get_nodes_from_documents(documents=documents)
print('doc:', documents)
print('22'.center(20, '='))
print('base_node:', base_node)
print('22'.center(20, '='))

# 根据节点，构建向量索引
index = VectorStoreIndex(nodes=base_node)
# 将索引持久化存储到本地向量数据库
index.storage_context.persist()

# 创建一个查询引擎，这个引擎可以接收查询并返回相关文档的响应
query_engine = index.as_query_engine()
rsp = query_engine.query('你叫什么名字？')
print('rsp:', rsp)
